
# Copyright (c) [2025] [Lent]
# [Deployment of Reinforcement Learning Navigation Models in ROS] is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan
# PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
# KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
# NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.


import torch
import torch.nn as nn
import torch.nn.functional as fu
import numpy as np

# 选择设备: cuda:0 或 cpu
device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
print(f"using device: {device}")


class ReplayBuffer:
    def __init__(self, capacity, obs_dim, state_dim, action_dim, batch_size):
        self.capacity = capacity
        self.obs_cap = np.empty((self.capacity, obs_dim))
        self.next_obs_cap = np.empty((self.capacity, obs_dim))
        self.state_cap = np.empty((self.capacity, state_dim))
        self.next_state_cap = np.empty((self.capacity, state_dim))
        self.action_cap = np.empty((self.capacity, action_dim))
        self.reward_cap = np.empty((self.capacity, 1))
        self.done_cap = np.empty((self.capacity, 1), dtype=bool)

        self.batch_size = batch_size
        self.current = 0

    def add_memo(self, obs, next_obs, state, next_state, action, reward, done):
        self.obs_cap[self.current] = obs
        self.next_obs_cap[self.current] = next_obs
        self.state_cap[self.current] = state
        self.next_state_cap[self.current] = next_state
        self.action_cap[self.current] = action
        self.reward_cap[self.current] = reward
        self.done_cap[self.current] = done

        self.current = (self.current + 1) % self.capacity

    def sample(self, idxes):
        obs = self.obs_cap[idxes]
        next_obs = self.next_obs_cap[idxes]
        state = self.state_cap[idxes]
        next_state = self.next_state_cap[idxes]
        action = self.action_cap[idxes]
        reward = self.reward_cap[idxes]
        done = self.done_cap[idxes]

        return obs, next_obs, state, next_state, action, reward, done

class Actor(nn.Module):
    def __init__(self, lr_actor, input_dims, fc1_dims, fc2_dims, action_dim):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(input_dims, fc1_dims)
        self.fc2 = nn.Linear(fc1_dims, fc2_dims)
        self.fc3 = nn.Linear(fc2_dims, fc2_dims)

        # 线速度输出层，使用sigmoid激活函数约束到[0,1]
        self.linear_velocity = nn.Linear(fc2_dims, 1)
        # 角速度输出层，使用tanh激活函数约束到[-1,1]
        self.angular_velocity = nn.Linear(fc2_dims, 1)

        nn.init.kaiming_uniform_(self.fc1.weight, a=0.01)  # 与LeakyReLU的0.01斜率匹配
        nn.init.kaiming_uniform_(self.fc2.weight, a=0.01)

        self.optimizer = torch.optim.Adam(self.parameters(), lr=lr_actor)

    def forward(self, state):
        x = fu.leaky_relu(self.fc1(state), negative_slope=0.01)
        x = fu.leaky_relu(self.fc2(x), negative_slope=0.01)
        x = fu.leaky_relu(self.fc3(x), negative_slope=0.01)
        # 并行输出线速度和角速度
        lin_vel = torch.sigmoid(self.linear_velocity(x))  # [0,1]范围
        ang_vel = torch.tanh(self.angular_velocity(x))  # [-1,1]范围
        # print(f"线速度为{lin_vel}角速度为：{ang_vel}")

        return torch.cat([lin_vel, ang_vel], dim=1)

    def save_checkpoint(self, checkpoint_file):
        torch.save(self.state_dict(), checkpoint_file)

    def load_checkpoint(self, checkpoint_file):
        self.load_state_dict(torch.load(checkpoint_file))


class Agent:
    def __init__(self, memo_size, obs_dim, state_dim, n_agent, action_dim,
                 alpha, beta, fc1_dims, fc2_dims, gamma, tau, batch_size):
        self.gamma = gamma
        self.tau = tau
        self.action_dim = action_dim

        self.actor = Actor(lr_actor=alpha, input_dims=obs_dim, fc1_dims=fc1_dims, fc2_dims=fc2_dims,
                           action_dim=action_dim).to(device)
        self.target_actor = Actor(lr_actor=alpha, input_dims=obs_dim, fc1_dims=fc1_dims, fc2_dims=fc2_dims,
                                  action_dim=action_dim).to(device)
        self.replay_buffer = ReplayBuffer(capacity=memo_size, obs_dim=obs_dim, state_dim=state_dim,
                                          action_dim=action_dim, batch_size=batch_size)

    def get_action(self, obs):
        single_obs = torch.tensor(obs, dtype=torch.float).unsqueeze(0).to(device)
        action_values = self.actor.forward(single_obs)

        # 添加高斯噪声探索（线速度噪声幅度0.05，角速度噪声幅度0.1）///仅在训练时启用！！！
        # noise = torch.randn_like(action_values) * torch.tensor([0.05, 0.1]).to(device)
        # noisy_actions = action_values + noise
        
        noisy_actions = action_values

        # 应用激活函数约束输出范围
        lin_vel = torch.clamp(noisy_actions[:, 0], 0.0, 1.0)  # 硬性截断
        ang_vel = torch.clamp(noisy_actions[:, 1], -1.0, 1.0)
        # 组合动作并转换为numpy数组
        action = torch.cat([lin_vel, ang_vel]).detach().cpu().numpy()
        return action

    def save_model(self, filename):
        self.actor.save_checkpoint(filename)
        self.target_actor.save_checkpoint(filename)
        self.critic.save_checkpoint(filename)
        self.target_critic.save_checkpoint(filename)

    def load_model(self, filename):
        self.actor.load_checkpoint(filename)
        self.target_actor.load_checkpoint(filename)
        self.critic.load_checkpoint(filename)
        self.target_critic.load_checkpoint(filename)
